import pandas as pd
import numpy as np


class GroupModeFiller:
    """
    用分组众数填充缺失值，如果测试数据没有分组列，则使用全局众数

    参数:
    group_col: 作为分组依据的列名（如果测试数据没有该列，则用全局众数）
    fill_cols: 需要填充的列名列表
    """

    def __init__(self, group_col, fill_cols):
        self.group_col = group_col
        self.fill_cols = fill_cols
        self.mode_dict = {}  # 存储 {列名: {分组1: 众数, 分组2: 众数, ...}}
        self.global_modes = {}  # 存储全局众数（用于测试数据没有 group_col 时）

    def fit(self, df):
        """从训练数据学习分组众数和全局众数"""
        self.mode_dict = {}
        self.global_modes = {}

        for col in self.fill_cols:
            # 计算每个分组的众数
            group_modes = df.groupby(self.group_col)[col].apply(
                lambda x: x.mode()[0] if not x.mode().empty else np.nan
            )
            self.mode_dict[col] = group_modes.to_dict()

            # 计算全局众数（用于测试数据没有 group_col 时）
            global_mode = df[col].mode()[0] if not df[col].mode().empty else np.nan
            self.global_modes[col] = global_mode
        return self

    def transform(self, df):
        """用训练数据的众数填充测试数据"""
        df_filled = df.copy()

        # 如果测试数据没有 group_col，直接用全局众数填充
        if self.group_col not in df_filled.columns:
            for col in self.fill_cols:
                df_filled[col] = df_filled[col].fillna(self.global_modes[col])
            return df_filled

        # 否则，按分组填充
        for col in self.fill_cols:
            # 用分组众数填充
            temp_col = df_filled[self.group_col].map(self.mode_dict[col])
            df_filled[col] = df_filled[col].fillna(temp_col)

            # 如果仍有 NaN（如新分组），用全局众数填充
            if df_filled[col].isna().any():
                df_filled[col] = df_filled[col].fillna(self.global_modes[col])

        return df_filled

    def fit_transform(self, df_train, df_test=None):
        """学习训练数据众数，并填充训练集（和测试集）"""
        self.fit(df_train)
        if df_test is not None:
            return self.transform(df_train), self.transform(df_test)
        return self.transform(df_train)